# 搜索二叉树的实现
# 节点类

class TreeNode:
    def __init__(self, key, val, left=None, right=None, parent=None):
        self.key = key
        self.payload = val
        self.leftChild = left
        self.rightChild = right
        self.parent = parent

    def hasLeftChild(self):
        return self.leftChild

    def hasRightChild(self):
        return self.rightChild

    def isLeftChild(self):
        return self.parent and self.parent.leftChild == self

    def isRightChild(self):
        return self.parent and self.parent.rightChild == self

    def isRoot(self):
        return not self.root

    def isLeaf(self):
        return not(self.leftChild or self.rightChild)

    def hasAnyChildren(self):
        return self.leftChild or self.rightChild

    def hasBothChildren(self):
        return self.rightChild and self.leftChild

    def replaceNodeData(self, key, val, lc, rc):
        self.key = key
        self.payload = val
        self.leftChild = lc
        self.rightChild = rc 
        if self.hasLeftChild():
            self.leftChild.parent = self
        if self.hasRightChild():
            self.rightChild.parent = self


    def __iter__(self):
        if self:
            if self.hasLeftChild():
                for elem in self.leftChild:
                    yield elem
            yield self.key 
            if self.hasRightChild():
                for elem in self.rightChild:
                    yield elem
    
    def findSuccessor(self):
        succ = None
        if self.hasRightChild():
            succ=self.rightChild.findMin()
        else:
            if self.parent:
                if self.isLeftChild():
                    succ=self.parent
                else:
                    self.parent.rightChild = None
                    succ=self.parent.findSuccessor()
                    self.parent.righChild = self
        return succ

    def findMin(self):
        current = self
        while current.hasLeftChild():
            current=current.leftChild
        return current

    def spliceOut(self):
        if self.isLeaf():
            if self.isLeftChild():
                self.parent.leftChild = None
            else:
                self.parent.rightChild = None
        elif self.hasAnyChildren():
            if self.hasLeftChild():
                if self.isLeftChild():
                    self.parent.leftChild = self.leftChild
                else:
                    self.parent.rightChild = self.leftChild
                self.leftChild.parent = self.parent
            else:
                if self.isLeftChild():
                    self.parent.leftChild = self.rightChild
                else:
                    self.parent.rightChild = self.rightChild
                self.rightChild.parent = self.parent

class BinarySearchTree:
    def __init__(self):
        self.root = None
        self.size = 0

    def length(self):
        return self.size

    def __len__(self):
        return self.size

    def __iter__(self):
        return self.root.__iter__()

    # 插入一个节点
    def put(self, key, val):
        if self.root:
            self._put(key, val, self.root)
        else:
            self.root = TreeNode(key,val)
        self.size += 1

    def _put(self, key, val, currentNode):
        if key < currentNode.key:
            if currentNode.hasLeftChild():
                self._put(key, val, currentNode.leftChild)
            else:
                currentNode.leftChild = TreeNode(key, val, parent=currentNode)
        else:
            if currentNode.hasRightChild():
                self._put(key, val, currentNode.rightChild)
            else:
                currentNode.rightChild = TreeNode(key, val, parent=currentNode)

    def __setitem__(self, k, v):
        self.put(k, v)

    # 查询节点:
    # 根据key值返回对于的value
    def get(self, key):
        if self.root:
            res = self._get(key, self.root)
            if res:
                return res.payload
            else:
                return None 
        else:
            return None

    # 返回待查找节点
    def _get(self, key, currentNode):
        if not currentNode:
            return None
        elif currentNode.key==key:
            return currentNode
        elif key<currentNode.key:
            return self._get(key, currentNode.leftChild)
        else:
            return self._get(key, currentNode.rightChild)

    def __getitem__(self, key):
        return self.get(key)

    def __contains__(self, key):
        if self._get(key, self.root):
            return True
        else:
            return False

    # 删除节点
    def delete(self, key):
        if self.size > 1:
            nodeToRemove = self._get(key, self.root)
            if nodeToRemove:
                self.remove(nodeToRemove)
                self.size -= 1
            else:
                raise KeyError("Error, key not in tree")
        elif self.size == 1 and self.root.key == key:
            self.root = None
            self.size -= 1
        else:
            raise KeyError("Error, key not in tree")

    def __delitem__(self, key):
        self.delete(key)


    def remove(self, currentNode):
        # 1 删除叶子节点
        if currentNode.isLeaf():
            # 1.1 删除左子节点(叶子)
            if currentNode==currentNode.parent.leftChild:
                currentNode.parent.leftChild = None
            # 1.2 删除右子节点(叶子)
            else:
                currentNode.parent.rightChild = None 
        # 3 删除含有两个子树的节点
        elif currentNode.hasBothChildren():
            print("here")
            succ = currentNode.findSuccessor()
            print(succ.key)
            succ.spliceOut()

            currentNode.key=succ.key
            currentNode.payload=succ.payload
        # 2 删除带有一个子树的节点
        # elif currentNode.hasAnyChildren():
        else:
            # 2.1 删除带有左子树的节点
            if currentNode.hasLeftChild():
                # 2.1.1 删除节点是父节点的左子树
                if currentNode.isLeftChild():
                    currentNode.leftChild.parent=currentNode.parent
                    currentNode.parent.leftChild = currentNode.leftChild
                # 2.1.2 删除节点是父节点的右子树
                elif currentNode.isRightChild():
                    currentNode.leftChild.parent = currentNode.parent
                    currentNode.parent.rightChild = currentNode.leftChild
                # 2.1.3 删除节点是根节点
                else:
                    currentNode.replaceNodeData(currentNode.leftChild.key,
                        currentNode.leftChild.payload,
                        currentNode.leftChild.leftChild,
                        currentNode.leftChild.rightChild)
            # 2.2 删除带有右子树的节点
            else:
                if currentNode.isLeftChild():
                    currentNode.rightChild.parent = currentNode.parent
                    currentNode.parent.leftChild = currentNode.rightChild
                elif currentNode.isRightChild():
                    currentNode.rightChild.parent = currentNode.parent
                    currentNode.parent.rightChild = currentNode.rightChild
                else:
                    currentNode.replaceNodeData(currentNode.rightChild.key,
                        currentNode.rightChild.payload,
                        currentNode.rightChild.leftChild,
                        currentNode.rightChild.rightChild)






# AVL树
# BST树的改进, 目的是为了保持BST树为平衡状态
class AVLNode(TreeNode):
    def __init__(self, key, val,left, right, parent):
        super().__init__(key, val, left=left, right=right, parent=parent)
        self.balanceFactor = 0

class AVLTree(BinarySearchTree):
    def _put(self, key, val, currentNode):
        if key < currentNode.key:
            if currentNode.hasLeftChild():
                self._put(key, val, currentNode.leftChild)
            else:
                currentNode.leftChild = AVLNode(key, val, parent=currentNode)
                self.updateBalance(currentNode.leftChild)
        else:
            if currentNode.hasRightChild():
                self._put(key, val, currentNode.rightChild)
            else:
                currentNode.rightChild = AVLNode(key, val, parent=currentNode)
                self.updateBalance(currentNode.rightChild)
        
    def updateBanlance(self, node):
        if node.balanceFactor > 1 or node.balanceFactor < -1:
            self.rebalance(node)
            return 

        if node.parent != None:
            if node.isLeftChild():
                node.parent.balanceFactor += 1
            elif node.isRightChild():
                node.parent.balanceFactor -= 1
            if node.parent.balanceFactor !=0:
                self.updateBanlance(node.parent)


    def rebalance(self, node):
        if node.balanceFactor < 0:
            if node.rightChild.balanceFactor > 0:
                self.rotateRight(node.rightChild)
                self.rotateLeft(node)
            else:
                self.rotateLeft(node)
        elif node.balanceFactor > 0:
            if node.leftChild.balanceFactor<0:
                self.rotateLeft(node.leftChild)
                self.rotateRight(node)
            else:
                self.rotateRight(node)

    def rotateLeft(self,rotRoot):
        newRoot = rotRoot.rightChild
        rotRoot.rightChild = newRoot.leftChild
        if newRoot.leftChild != None:
            newRoot.leftChild.parent = rotRoot
        newRoot.parent = rotRoot.parent
        if rotRoot.isRoot():
            self.root = newRoot
        else:
            if rotRoot.isLeftChild():
                    rotRoot.parent.leftChild = newRoot
            else:
                rotRoot.parent.rightChild = newRoot
        newRoot.leftChild = rotRoot
        rotRoot.parent = newRoot
        rotRoot.balanceFactor = rotRoot.balanceFactor + 1 - min(newRoot.balanceFactor, 0)
        newRoot.balanceFactor = newRoot.balanceFactor + 1 + max(rotRoot.balanceFactor, 0)




bst = BinarySearchTree()

bst[50]="red"
bst.put(44, "blue")
bst.put(70, "yelllow")
bst.put(60, 'kk')
bst[40]="blueleft"
bst[90]="yellowright"
bst[88]="yellowrl"
bst[91]='yellowlr'
bst[89]="test"
print(len(bst))

print(bst[44])
print(50 in bst)

# node=bst._get(44, bst.root)
# print(node.isLeftChild())
print("删除前:")
for key in  bst:
    print(key, bst[key])
# del bst[60]
# del bst[44]
# print("删除后:")
# for key in  bst:
#     print(key, bst[key])

del bst[70]

print("删除后:")
for key in  bst:
    print(key, bst[key])